{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b7549a27-bc4a-4609-bb25-cc7d95cf8c23",
   "metadata": {},
   "source": [
    "# Preprocessing Chat History with `TransformMessages`\n",
    "\n",
    "## Introduction\n",
    "This notebook illustrates how to use `TransforMessages` give any `ConversableAgent` the ability to handle long contexts, sensitive data, and more.\n",
    "\n",
    "````{=mdx}\n",
    ":::info Requirements\n",
    "Install `pyautogen`:\n",
    "```bash\n",
    "pip install pyautogen\n",
    "```\n",
    "\n",
    "For more information, please refer to the [installation guide](/docs/installation/).\n",
    ":::\n",
    "````"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "47773f79-c0fd-4993-bc6e-3d1a57690118",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pprint\n",
    "import copy\n",
    "import re\n",
    "\n",
    "import autogen\n",
    "from autogen.agentchat.contrib.capabilities import transform_messages, transforms\n",
    "from typing import Dict, List"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9f09246b-a7d0-4238-b62c-1e72c7d815b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "config_list = autogen.config_list_from_json(\n",
    "    env_or_file=\"OAI_CONFIG_LIST\",\n",
    ")\n",
    "# Define your llm config\n",
    "llm_config = {\"config_list\": config_list}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea68962a-048d-42e9-9fca-cd944c56184d",
   "metadata": {},
   "source": [
    "````{=mdx}\n",
    ":::tip\n",
    "Learn more about configuring LLMs for agents [here](/docs/topics/llm_configuration).\n",
    ":::\n",
    "````"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "84d0e5ad-8b35-4b30-847e-4723e9c76f7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define your agent; the user proxy and an assistant\n",
    "assistant = autogen.AssistantAgent(\n",
    "    \"assistant\",\n",
    "    llm_config=llm_config,\n",
    ")\n",
    "user_proxy = autogen.UserProxyAgent(\n",
    "    \"user_proxy\",\n",
    "    human_input_mode=\"NEVER\",\n",
    "    is_termination_msg=lambda x: \"TERMINATE\" in x.get(\"content\", \"\"),\n",
    "    max_consecutive_auto_reply=10,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "180aa953-45be-469a-a94f-0ed0b4ef5ddf",
   "metadata": {},
   "source": [
    "## Handling Long Contexts\n",
    "\n",
    "Imagine a scenario where the LLM generates an extensive amount of text, surpassing the token limit imposed by your API provider. To address this issue, you can leverage `TransformMessages` along with its constituent transformations, `MessageHistoryLimiter` and `MessageTokenLimiter`.\n",
    "\n",
    "- `MessageHistoryLimiter`: You can restrict the total number of messages considered as context history. This transform is particularly useful when you want to limit the conversational context to a specific number of recent messages, ensuring efficient processing and response generation.\n",
    "- `MessageTokenLimiter`: Enables you to cap the total number of tokens, either on a per-message basis or across the entire context history (or both). This transformation is invaluable when you need to adhere to strict token limits imposed by your API provider, preventing unnecessary costs or errors caused by exceeding the allowed token count."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "34b943a2-ec58-41bc-a449-d9118c4bbdea",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Limit the message history to the 3 most recent messages\n",
    "max_msg_transfrom = transforms.MessageHistoryLimiter(max_messages=3)\n",
    "\n",
    "# Limit the token limit per message to 10 tokens\n",
    "token_limit_transform = transforms.MessageTokenLimiter(max_tokens_per_message=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "679c1026-4e1b-4c07-85cc-86594cc0b87b",
   "metadata": {},
   "source": [
    "## Example 1: Limiting number of messages\n",
    "Let's take a look at how these transformations will effect the messages. Below we see that by applying the `MessageHistoryLimiter`, we can see that we limited the context history to the 3 most recent messages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "61a2ead4-5f8b-4108-b1f0-3b51b41e2231",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'content': 'how', 'role': 'user'},\n",
      " {'content': [{'text': 'are you doing?', 'type': 'text'}], 'role': 'assistant'},\n",
      " {'content': 'very very very very very very long string', 'role': 'user'}]\n"
     ]
    }
   ],
   "source": [
    "messages = [\n",
    "    {\"role\": \"user\", \"content\": \"hello\"},\n",
    "    {\"role\": \"assistant\", \"content\": [{\"type\": \"text\", \"text\": \"there\"}]},\n",
    "    {\"role\": \"user\", \"content\": \"how\"},\n",
    "    {\"role\": \"assistant\", \"content\": [{\"type\": \"text\", \"text\": \"are you doing?\"}]},\n",
    "    {\"role\": \"user\", \"content\": \"very very very very very very long string\"},\n",
    "]\n",
    "\n",
    "processed_messages = max_msg_transfrom.apply_transform(copy.deepcopy(messages))\n",
    "pprint.pprint(processed_messages)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "610739af-b812-404e-82d2-b3ed796b8b6c",
   "metadata": {},
   "source": [
    "## Example 2: Limiting number of tokens\n",
    "\n",
    "Now let's test limiting the number of tokens in messages. We can see that we can limit the number of tokens to 3, which is equivalent to 3 words in this instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "739dd260-fa95-4e5d-ae84-9cb7f40de975",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mTruncated 6 tokens. Tokens reduced from 15 to 9\u001b[0m\n",
      "[{'content': 'hello', 'role': 'user'},\n",
      " {'content': [{'text': 'there', 'type': 'text'}], 'role': 'assistant'},\n",
      " {'content': 'how', 'role': 'user'},\n",
      " {'content': [{'text': 'are you doing', 'type': 'text'}], 'role': 'assistant'},\n",
      " {'content': 'very very very', 'role': 'user'}]\n"
     ]
    }
   ],
   "source": [
    "processed_messages = token_limit_transform.apply_transform(copy.deepcopy(messages))\n",
    "\n",
    "pprint.pprint(processed_messages)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35fa2844-bd83-42ac-8275-959f093b7bc7",
   "metadata": {},
   "source": [
    "## Example 3: Combining transformations\n",
    "\n",
    "Let's test these transforms with agents (the upcoming test is replicated from the agentchat_capability_long_context_handling notebook). We will see that the agent without the capability to handle long context will result in an error, while the agent with that capability will have no issues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "80e53623-2830-41b7-8ae2-bf3668071657",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33muser_proxy\u001b[0m (to assistant):\n",
      "\n",
      "plot and save a graph of x^2 from -10 to 10\n",
      "\n",
      "--------------------------------------------------------------------------------\n",
      "Encountered an error with the base assistant\n",
      "Error code: 429 - {'error': {'message': 'Request too large for gpt-3.5-turbo in organization org-U58JZBsXUVAJPlx2MtPYmdx1 on tokens per min (TPM): Limit 60000, Requested 1252546. The input or output tokens must be reduced in order to run successfully. Visit https://platform.openai.com/account/rate-limits to learn more.', 'type': 'tokens', 'param': None, 'code': 'rate_limit_exceeded'}}\n",
      "\n",
      "\n",
      "\n",
      "\u001b[33muser_proxy\u001b[0m (to assistant):\n",
      "\n",
      "plot and save a graph of x^2 from -10 to 10\n",
      "\n",
      "--------------------------------------------------------------------------------\n",
      "\u001b[33mTruncated 3804 tokens. Tokens reduced from 4019 to 215\u001b[0m\n",
      "\u001b[33massistant\u001b[0m (to user_proxy):\n",
      "\n",
      "To plot the graph of \\( x^2 \\) from -10 to 10 and save it, we can use Python with the matplotlib library. Here is the code to achieve this:\n",
      "\n",
      "```python\n",
      "# filename: plot_graph.py\n",
      "import matplotlib.pyplot as plt\n",
      "import numpy as np\n",
      "\n",
      "x = np.linspace(-10, 10, 100)\n",
      "y = x**2\n",
      "\n",
      "plt.plot(x, y)\n",
      "plt.xlabel('x')\n",
      "plt.ylabel('x^2')\n",
      "plt.title('Graph of x^2')\n",
      "plt.grid(True)\n",
      "plt.savefig('x_squared_graph.png')\n",
      "plt.show()\n",
      "```\n",
      "\n",
      "After executing this code, you should see the graph of \\( x^2 \\) displayed and saved as `x_squared_graph.png`.\n",
      "\n",
      "Please make sure you have matplotlib installed. If not, you can install it using pip:\n",
      "\n",
      "```sh\n",
      "pip install matplotlib\n",
      "```\n",
      "\n",
      "Go ahead and execute the Python script provided above to plot and save the graph of \\( x^2 \\). Let me know if you encounter any issues.\n",
      "\n",
      "--------------------------------------------------------------------------------\n",
      "\u001b[31m\n",
      ">>>>>>>> EXECUTING CODE BLOCK 0 (inferred language is python)...\u001b[0m\n",
      "\u001b[31m\n",
      ">>>>>>>> EXECUTING CODE BLOCK 1 (inferred language is sh)...\u001b[0m\n",
      "\u001b[33muser_proxy\u001b[0m (to assistant):\n",
      "\n",
      "exitcode: 0 (execution succeeded)\n",
      "Code output: \n",
      "Figure(640x480)\n",
      "\n",
      "Requirement already satisfied: matplotlib in /home/wael/workspaces/autogen/.venv/lib/python3.11/site-packages (3.8.2)\n",
      "Requirement already satisfied: contourpy>=1.0.1 in /home/wael/workspaces/autogen/.venv/lib/python3.11/site-packages (from matplotlib) (1.2.0)\n",
      "Requirement already satisfied: cycler>=0.10 in /home/wael/workspaces/autogen/.venv/lib/python3.11/site-packages (from matplotlib) (0.12.1)\n",
      "Requirement already satisfied: fonttools>=4.22.0 in /home/wael/workspaces/autogen/.venv/lib/python3.11/site-packages (from matplotlib) (4.48.1)\n",
      "Requirement already satisfied: kiwisolver>=1.3.1 in /home/wael/workspaces/autogen/.venv/lib/python3.11/site-packages (from matplotlib) (1.4.5)\n",
      "Requirement already satisfied: numpy<2,>=1.21 in /home/wael/workspaces/autogen/.venv/lib/python3.11/site-packages (from matplotlib) (1.26.4)\n",
      "Requirement already satisfied: packaging>=20.0 in /home/wael/workspaces/autogen/.venv/lib/python3.11/site-packages (from matplotlib) (23.2)\n",
      "Requirement already satisfied: pillow>=8 in /home/wael/workspaces/autogen/.venv/lib/python3.11/site-packages (from matplotlib) (10.2.0)\n",
      "Requirement already satisfied: pyparsing>=2.3.1 in /home/wael/workspaces/autogen/.venv/lib/python3.11/site-packages (from matplotlib) (3.1.1)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /home/wael/workspaces/autogen/.venv/lib/python3.11/site-packages (from matplotlib) (2.8.2)\n",
      "Requirement already satisfied: six>=1.5 in /home/wael/workspaces/autogen/.venv/lib/python3.11/site-packages (from python-dateutil>=2.7->matplotlib) (1.16.0)\n",
      "\n",
      "\n",
      "--------------------------------------------------------------------------------\n",
      "\u001b[33mTruncated 3435 tokens. Tokens reduced from 3700 to 265\u001b[0m\n",
      "\u001b[33massistant\u001b[0m (to user_proxy):\n",
      "\n",
      "The graph has been successfully created and saved. You can find the graph as a file named \"x_squared_plot.png\" in the directory where you ran the script. You can open and view this file to see the plotted graph of \\(x^2\\) from -10 to 10.\n",
      "\n",
      "TERMINATE\n",
      "\n",
      "--------------------------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "llm_config = {\n",
    "    \"config_list\": [{\"model\": \"gpt-3.5-turbo\", \"api_key\": os.environ.get(\"OPENAI_API_KEY\")}],\n",
    "}\n",
    "\n",
    "assistant_base = autogen.AssistantAgent(\n",
    "    \"assistant\",\n",
    "    llm_config=llm_config,\n",
    ")\n",
    "\n",
    "assistant_with_context_handling = autogen.AssistantAgent(\n",
    "    \"assistant\",\n",
    "    llm_config=llm_config,\n",
    ")\n",
    "# suppose this capability is not available\n",
    "context_handling = transform_messages.TransformMessages(\n",
    "    transforms=[\n",
    "        transforms.MessageHistoryLimiter(max_messages=10),\n",
    "        transforms.MessageTokenLimiter(max_tokens=1000, max_tokens_per_message=50),\n",
    "    ]\n",
    ")\n",
    "\n",
    "context_handling.add_to_agent(assistant_with_context_handling)\n",
    "\n",
    "user_proxy = autogen.UserProxyAgent(\n",
    "    \"user_proxy\",\n",
    "    human_input_mode=\"NEVER\",\n",
    "    is_termination_msg=lambda x: \"TERMINATE\" in x.get(\"content\", \"\"),\n",
    "    code_execution_config={\n",
    "        \"work_dir\": \"coding\",\n",
    "        \"use_docker\": False,\n",
    "    },\n",
    "    max_consecutive_auto_reply=2,\n",
    ")\n",
    "\n",
    "# suppose the chat history is large\n",
    "# Create a very long chat history that is bound to cause a crash\n",
    "# for gpt 3.5\n",
    "for i in range(1000):\n",
    "    # define a fake, very long messages\n",
    "    assitant_msg = {\"role\": \"assistant\", \"content\": \"test \" * 1000}\n",
    "    user_msg = {\"role\": \"user\", \"content\": \"\"}\n",
    "\n",
    "    assistant_base.send(assitant_msg, user_proxy, request_reply=False, silent=True)\n",
    "    assistant_with_context_handling.send(assitant_msg, user_proxy, request_reply=False, silent=True)\n",
    "    user_proxy.send(user_msg, assistant_base, request_reply=False, silent=True)\n",
    "    user_proxy.send(user_msg, assistant_with_context_handling, request_reply=False, silent=True)\n",
    "\n",
    "try:\n",
    "    user_proxy.initiate_chat(assistant_base, message=\"plot and save a graph of x^2 from -10 to 10\", clear_history=False)\n",
    "except Exception as e:\n",
    "    print(\"Encountered an error with the base assistant\")\n",
    "    print(e)\n",
    "    print(\"\\n\\n\")\n",
    "\n",
    "try:\n",
    "    user_proxy.initiate_chat(\n",
    "        assistant_with_context_handling, message=\"plot and save a graph of x^2 from -10 to 10\", clear_history=False\n",
    "    )\n",
    "except Exception as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e380678-a923-43cb-91b1-f9c9e8deede2",
   "metadata": {},
   "source": [
    "## Handling Sensitive Data\n",
    "\n",
    "You can use the `MessageTransform` protocol to create custom message transformations that redact sensitive data from the chat history. This is particularly useful when you want to ensure that sensitive information, such as API keys, passwords, or personal data, is not exposed in the chat history or logs.\n",
    "\n",
    "Now, we will create a custom message transform to detect any OpenAI API key and redact it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "74429344-3c0a-4057-aba3-27358fbf059c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The transform must adhere to transform_messages.MessageTransform protocol.\n",
    "class MessageRedact:\n",
    "    def __init__(self):\n",
    "        self._openai_key_pattern = r\"sk-([a-zA-Z0-9]{48})\"\n",
    "        self._replacement_string = \"REDACTED\"\n",
    "\n",
    "    def apply_transform(self, messages: List[Dict]) -> List[Dict]:\n",
    "        temp_messages = copy.deepcopy(messages)\n",
    "\n",
    "        for message in temp_messages:\n",
    "            if isinstance(message[\"content\"], str):\n",
    "                message[\"content\"] = re.sub(self._openai_key_pattern, self._replacement_string, message[\"content\"])\n",
    "            elif isinstance(message[\"content\"], list):\n",
    "                for item in message[\"content\"]:\n",
    "                    if item[\"type\"] == \"text\":\n",
    "                        item[\"text\"] = re.sub(self._openai_key_pattern, self._replacement_string, item[\"text\"])\n",
    "        return temp_messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8a79c0b4-5ff8-49c5-b8a6-c54ca4c7cca2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33muser_proxy\u001b[0m (to assistant):\n",
      "\n",
      "What are the two API keys that I just provided\n",
      "\n",
      "--------------------------------------------------------------------------------\n",
      "\u001b[33massistant\u001b[0m (to user_proxy):\n",
      "\n",
      "To retrieve the two API keys you provided, I will display them individually in the output. \n",
      "\n",
      "Here is the first API key:\n",
      "```python\n",
      "# Display the first API key\n",
      "print(\"API key 1 =\", \"REDACTED\")\n",
      "```\n",
      "\n",
      "Here is the second API key:\n",
      "```python\n",
      "# Display the second API key\n",
      "print(\"API key 2 =\", \"REDACTED\")\n",
      "```\n",
      "\n",
      "Please run the code snippets to see the API keys. After that, I will mark this task as complete.\n",
      "\n",
      "--------------------------------------------------------------------------------\n",
      "\u001b[31m\n",
      ">>>>>>>> EXECUTING CODE BLOCK 0 (inferred language is python)...\u001b[0m\n",
      "\u001b[31m\n",
      ">>>>>>>> EXECUTING CODE BLOCK 1 (inferred language is python)...\u001b[0m\n",
      "\u001b[33muser_proxy\u001b[0m (to assistant):\n",
      "\n",
      "exitcode: 0 (execution succeeded)\n",
      "Code output: \n",
      "API key 1 = REDACTED\n",
      "\n",
      "API key 2 = REDACTED\n",
      "\n",
      "\n",
      "--------------------------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "assistant_with_redact = autogen.AssistantAgent(\n",
    "    \"assistant\",\n",
    "    llm_config=llm_config,\n",
    "    max_consecutive_auto_reply=1,\n",
    ")\n",
    "# suppose this capability is not available\n",
    "redact_handling = transform_messages.TransformMessages(transforms=[MessageRedact()])\n",
    "\n",
    "redact_handling.add_to_agent(assistant_with_redact)\n",
    "\n",
    "user_proxy = autogen.UserProxyAgent(\n",
    "    \"user_proxy\",\n",
    "    human_input_mode=\"NEVER\",\n",
    "    max_consecutive_auto_reply=1,\n",
    ")\n",
    "\n",
    "messages = [\n",
    "    {\"content\": \"api key 1 = sk-7nwt00xv6fuegfu3gnwmhrgxvuc1cyrhxcq1quur9zvf05fy\"},  # Don't worry, randomly generated\n",
    "    {\"content\": [{\"type\": \"text\", \"text\": \"API key 2 = sk-9wi0gf1j2rz6utaqd3ww3o6c1h1n28wviypk7bd81wlj95an\"}]},\n",
    "]\n",
    "\n",
    "for message in messages:\n",
    "    user_proxy.send(message, assistant_with_redact, request_reply=False, silent=True)\n",
    "\n",
    "result = user_proxy.initiate_chat(\n",
    "    assistant_with_redact, message=\"What are the two API keys that I just provided\", clear_history=False\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "Preprocessing chat history with `TransformMessages`",
   "tags": [
    "long context handling",
    "capability"
   ]
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
